"""模型管理器模块"""

import os
from typing import Dict, Optional, List
from models.yolov8_model import YOLOv8Model
from config import get_config

class ModelManager:
    """模型管理器，负责管理和切换多个YOLO模型"""
    
    def __init__(self):
        """初始化模型管理器"""
        self.models: Dict[str, YOLOv8Model] = {}
        self.active_model_name: Optional[str] = None
        self.device = get_config('performance.device', 'auto')
        
        # 从配置加载模型列表
        self._load_models_from_config()
    
    def _load_models_from_config(self) -> None:
        """从配置文件加载模型列表"""
        model_list = get_config('models.list', [])
        
        for model_config in model_list:
            try:
                model_name = model_config.get('name', 'default')
                model_path = model_config.get('path', '')
                
                if model_path:
                    self.load_model(model_name, model_path)
                    
                    # 如果配置了默认模型，则设置为活动模型
                    if model_config.get('default', False):
                        self.set_active_model(model_name)
                        
            except Exception as e:
                print(f"加载模型 {model_config.get('name', 'unknown')} 失败: {e}")
    
    def load_model(self, model_name: str, model_path: str) -> bool:
        """加载指定的模型
        
        Args:
            model_name: 模型名称
            model_path: 模型文件路径
            
        Returns:
            是否加载成功
        """
        try:
            # 检查模型文件是否存在
            if not os.path.isabs(model_path):
                base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
                model_path = os.path.join(base_dir, model_path)
            
            if not os.path.exists(model_path):
                print(f"模型文件不存在: {model_path}")
                return False
            
            # 创建模型实例
            model = YOLOv8Model(model_path=model_path, device=self.device)
            self.models[model_name] = model
            
            print(f"模型 '{model_name}' 已成功加载")
            
            # 如果是第一个加载的模型，自动设置为活动模型
            if self.active_model_name is None:
                self.set_active_model(model_name)
                
            return True
            
        except Exception as e:
            print(f"加载模型 '{model_name}' 失败: {e}")
            return False
    
    def unload_model(self, model_name: str) -> bool:
        """卸载指定的模型
        
        Args:
            model_name: 模型名称
            
        Returns:
            是否卸载成功
        """
        try:
            if model_name in self.models:
                del self.models[model_name]
                
                # 如果卸载的是活动模型，尝试设置另一个模型为活动模型
                if model_name == self.active_model_name:
                    self.active_model_name = next(iter(self.models.keys()), None)
                    if self.active_model_name:
                        print(f"自动切换活动模型为: {self.active_model_name}")
                    else:
                        print("没有可用的模型")
                
                print(f"模型 '{model_name}' 已成功卸载")
                return True
            else:
                print(f"模型 '{model_name}' 不存在")
                return False
                
        except Exception as e:
            print(f"卸载模型 '{model_name}' 失败: {e}")
            return False
    
    def set_active_model(self, model_name: str) -> bool:
        """设置活动模型
        
        Args:
            model_name: 模型名称
            
        Returns:
            是否设置成功
        """
        if model_name in self.models:
            self.active_model_name = model_name
            print(f"活动模型已切换为: {model_name}")
            return True
        else:
            print(f"模型 '{model_name}' 不存在")
            return False
    
    def get_active_model(self) -> Optional[YOLOv8Model]:
        """获取当前活动模型
        
        Returns:
            当前活动的YOLOv8Model实例，如果没有活动模型则返回None
        """
        if self.active_model_name and self.active_model_name in self.models:
            return self.models[self.active_model_name]
        return None
    
    def get_model(self, model_name: str) -> Optional[YOLOv8Model]:
        """获取指定名称的模型
        
        Args:
            model_name: 模型名称
            
        Returns:
            YOLOv8Model实例，如果模型不存在则返回None
        """
        return self.models.get(model_name)
    
    def list_models(self) -> List[str]:
        """列出所有已加载的模型
        
        Returns:
            模型名称列表
        """
        return list(self.models.keys())
    
    def get_active_model_name(self) -> Optional[str]:
        """获取当前活动模型的名称
        
        Returns:
            活动模型名称，如果没有活动模型则返回None
        """
        return self.active_model_name
    
    def set_device(self, device: str) -> None:
        """设置所有模型的运行设备
        
        Args:
            device: 设备类型，'auto'、'cuda'、'cpu'
        """
        self.device = device
        
        # 更新所有已加载模型的设备
        active_model_name = self.active_model_name
        self.models.clear()
        self.active_model_name = None
        
        # 重新加载所有模型
        self._load_models_from_config()
        
        # 恢复活动模型
        if active_model_name and active_model_name in self.models:
            self.set_active_model(active_model_name)
    
    def get_model_info(self, model_name: str) -> Dict:
        """获取模型信息
        
        Args:
            model_name: 模型名称
            
        Returns:
            模型信息字典
        """
        model = self.models.get(model_name)
        if model:
            return model.get_performance_metrics()
        return {}
    
    def reload_all_models(self) -> None:
        """重新加载所有模型"""
        active_model_name = self.active_model_name
        self.models.clear()
        self.active_model_name = None
        
        # 重新加载所有模型
        self._load_models_from_config()
        
        # 恢复活动模型
        if active_model_name and active_model_name in self.models:
            self.set_active_model(active_model_name)

# 创建全局模型管理器实例
model_manager = ModelManager()

# 导出常用方法
def get_active_model() -> Optional[YOLOv8Model]:
    """获取当前活动模型的快捷方法"""
    return model_manager.get_active_model()

def set_active_model(model_name: str) -> bool:
    """设置活动模型的快捷方法"""
    return model_manager.set_active_model(model_name)

def list_loaded_models() -> List[str]:
    """列出所有已加载模型的快捷方法"""
    return model_manager.list_models()